{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a9c0da19-4ccb-41cc-8e57-f8e5426d87e3",
   "metadata": {},
   "source": [
    "# Building CLIP Model for Zero-Shot Image Classification with TT-NN\n",
    "\n",
    "CLIP (Contrastive Language-Image Pre-Training) is a foundational multimodal AI model developed by OpenAI that learns visual concepts from natural language supervision. Unlike traditional computer vision models that are trained on fixed categories, CLIP can understand and classify images based on arbitrary text descriptions.\n",
    "\n",
    "## What CLIP Does\n",
    "\n",
    "CLIP bridges the gap between vision and language by associating images with their textual descriptions. CLIP consists of two main components:\n",
    "\n",
    "1. **Vision Encoder**: A Vision Transformer (ViT) that processes images and converts them into feature embeddings.\n",
    "2. **Text Encoder**: A Transformer that processes text descriptions and converts them into feature embeddings.\n",
    "\n",
    "![CLIP Diagram](https://media.githubusercontent.com/media/tenstorrent/tutorial-assets/refs/heads/main/media/clip_tutorial/CLIP.png)\n",
    "\n",
    "During inference, CLIP can:\n",
    "- **Zero-shot image classification**: Classify images into categories it has never explicitly seen during training by comparing image embeddings with text embeddings of category descriptions.\n",
    "- **Image-text similarity**: Measure how well an image matches a given text description.\n",
    "- **Content-based image retrieval**: Find images that best match a text query.\n",
    "\n",
    "\n",
    "\n",
    "In this tutorial, we will implement CLIP for image classification. CLIP classifies images using natural language prompts like \"a diagram\", \"a dog\", or \"a cat\". \n",
    "\n",
    "The following diagram illustrates CLIP-based image classification of a list of prompts, and an image. CLIP computes the probability that a prompt matches the image using two pipelines:\n",
    "\n",
    "- A **Vision Pipeline** - Includes an image pre-processing step (resize, center image, etc.) and the vision encoder.\n",
    "- A **Text Pipeline** - Includes a tokenizer to convert prompts into a tensor, and the text encoder.\n",
    "\n",
    "Once both pipelines have been performed, their results are combined through a matrix multiplication, which gives us the probability for each prompt.\n",
    "\n",
    "\n",
    "![CLIP Zero Shot Classification](https://media.githubusercontent.com/media/tenstorrent/tutorial-assets/refs/heads/main/media/clip_tutorial/clip_zero_shot_classification.png)\n",
    "\n",
    "\n",
    "We use OpenAI's pre-trained weights for the clip-vit-base-patch32 model and focus on inference."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afc6c99f-80a5-4002-a3e5-6b4029f21705",
   "metadata": {},
   "source": [
    "## Imports and Dependencies\n",
    "\n",
    "We start by importing the necessary libraries for our CLIP implementation:\n",
    "\n",
    "- **TT-NN**: Tenstorrent's neural network library for hardware-accelerated deep learning.\n",
    "- **Torch**: PyTorch for model loading and tensor pre-processing.\n",
    "- **Transformers**: Hugging Face library for downloading pre-trained models and tokenizing text prompts.\n",
    "- **PIL**: Python Imaging Library for image pre-processing.\n",
    "- **Torchvision**: Computer vision utilities for image preprocessing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3d82f03-2f2d-4460-92a0-c6021ff4244c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ttnn\n",
    "import torch\n",
    "from loguru import logger\n",
    "import re\n",
    "import math\n",
    "from PIL import Image\n",
    "from transformers import CLIPTokenizer, CLIPModel\n",
    "import requests\n",
    "from io import BytesIO\n",
    "import time\n",
    "import safetensors.torch\n",
    "import os\n",
    "\n",
    "from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize, InterpolationMode"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecbf6215-2106-433a-a5cc-ca653e171a8b",
   "metadata": {},
   "source": [
    "## TT-NN Device Management and Utility Functions\n",
    "\n",
    "Helper functions are defined as TT-NN management devices and simplify device initialization.\n",
    "\n",
    "**Note on L1 Small Size**: For our model, we use a custom L1 memory configuration.  The `l1_small_size` parameter (8192 bytes) sets aside a portion of the on-chip L1 memory for sliding-window operations such as convolutions.\n",
    "Without `l1_small_size=8192`, convolution kernel would fail due to a lack of available L1 space. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e0e772f-4de5-4612-b615-0bf1a8d295b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def open_ttnn():\n",
    "    \"\"\"Initialize TT-NN device with specified L1 cache size.\"\"\"\n",
    "    global device\n",
    "    device = ttnn.open_device(device_id=0, l1_small_size=8192)\n",
    "\n",
    "def close_ttnn():\n",
    "    \"\"\"Clean up and close the TT-NN device.\"\"\"\n",
    "    global device\n",
    "    if device is not None:\n",
    "        ttnn.close_device(device)\n",
    "\n",
    "def get_device():\n",
    "    \"\"\"Get the current TT-NN device handle.\"\"\"\n",
    "    global device\n",
    "    return device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c632f6ee-1653-4d83-943b-9efbfbfdb444",
   "metadata": {},
   "source": [
    "## Model Weight Conversion\n",
    "\n",
    "TT-NN uses a custom tensor format that converts torch to TT-NN and vice versa. We can write a function to convert the entire model (more specifically, its `state_dict`) from PyTorch tensors to TT-NN tensors, enabling pre-trained CLIP weights on TT hardware.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18bb9e44-9d89-4606-abe6-76d7e37c597f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_model_to_ttnn(state_dict):\n",
    "    \"\"\"\n",
    "    Convert a PyTorch model's state dictionary to TT-NN format.\n",
    "    \n",
    "    Args:\n",
    "        state_dict: PyTorch model state dictionary containing weights and biases\n",
    "        \n",
    "    Returns:\n",
    "        dict: State dictionary with tensors converted to TT-NN format\n",
    "    \"\"\"\n",
    "    ttnn_state_dict = {}\n",
    "    logger.info(f\"Converting model to TT-NN format\")\n",
    "\n",
    "    # Convert each tensor in the state dictionary to TT-NN format\n",
    "    for key, value in state_dict.items():\n",
    "        if isinstance(value, torch.Tensor):\n",
    "            # Convert PyTorch tensors to TT-NN tensors\n",
    "            ttnn_state_dict[key] = ttnn.from_torch(value, layout=ttnn.TILE_LAYOUT, device=get_device())\n",
    "            ttnn_state_dict[key] = ttnn.typecast(ttnn_state_dict[key], dtype=ttnn.bfloat16)\n",
    "        elif isinstance(value, torch.Size):\n",
    "            # Convert PyTorch Size objects to TT-NN Size objects\n",
    "            ttnn_state_dict[key] = ttnn.Size(value)\n",
    "\n",
    "    return ttnn_state_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1d1d9a7-9f4f-446f-9b3a-d5f002b2abdb",
   "metadata": {},
   "source": [
    "## Generic Transformer Implementation\n",
    "\n",
    "CLIP uses two types of transformers: a text transformer and a vision transformer. To maximize code reuse, we define a generic Transformer class that can be used for both modalities with the appropriate configuration.\n",
    "\n",
    "### Transformer Architecture\n",
    "\n",
    "The transformer models used by CLIP consist of multiple layers (residual blocks), each containing the following sub-operations in sequence:\n",
    "\n",
    "1. **Layer Normalization**: Normalizes inputs for stable inference and training.\n",
    "2. **Multi-Head Self-Attention**: \n",
    "   - **Text**: Uses causal masking to prevent execution of future tokens.\n",
    "   - **Vision**: Uses full attention across all image patches.\n",
    "3. **Layer Normalization**: Second normalization layer.\n",
    "4. **MLP (Multi-Layer Perceptron)**: Two linear layers with GELU activation (Linear → GELU → Linear).\n",
    "\n",
    "Each block uses residual connections, where the output of each sub-operation is added to its input, enabling deeper networks and better gradient flow. \n",
    "\n",
    "Each block uses pre-trained weights and biases. These weights and biases are provided at initialization through the `state_dict` dictionary. This dictionary contains TT-NN tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a7f12bd-eeea-4e39-a43b-7498ad7a9fc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention:\n",
    "    def __init__(self, state_dict, num_heads, attention_mask=None, prefix=\"\"):\n",
    "        self.attention_mask = attention_mask\n",
    "        self.prefix = prefix\n",
    "        self.num_heads = num_heads\n",
    "\n",
    "        # Scale factor for attention scores: 1/sqrt(head_dim) for numerical stability\n",
    "        self.scale = 1.0 / math.sqrt(num_heads)\n",
    "\n",
    "        self.q_proj_weight = state_dict[f\"{prefix}.q_proj.weight\"]\n",
    "        self.q_proj_bias = state_dict[f\"{prefix}.q_proj.bias\"]\n",
    "        self.k_proj_weight = state_dict[f\"{prefix}.k_proj.weight\"]\n",
    "        self.k_proj_bias = state_dict[f\"{prefix}.k_proj.bias\"]\n",
    "        self.v_proj_weight = state_dict[f\"{prefix}.v_proj.weight\"]\n",
    "        self.v_proj_bias = state_dict[f\"{prefix}.v_proj.bias\"]\n",
    "        self.out_proj_weight = state_dict[f\"{prefix}.out_proj.weight\"]\n",
    "        self.out_proj_bias = state_dict[f\"{prefix}.out_proj.bias\"]\n",
    "\n",
    "\n",
    "    def forward(self, hidden_states):\n",
    "        sequence_size, batch_size, hidden_size = hidden_states.shape\n",
    "\n",
    "        head_size = hidden_size // self.num_heads\n",
    "\n",
    "        # Compute Q, K, V projections from input hidden states\n",
    "        # Each projection: [sequence_size, batch_size, hidden_size] -> [sequence_size, batch_size, hidden_size]\n",
    "        q = ttnn.linear(hidden_states, self.q_proj_weight, bias=self.q_proj_bias, transpose_b=True)\n",
    "        k = ttnn.linear(hidden_states, self.k_proj_weight, bias=self.k_proj_bias, transpose_b=True)\n",
    "        v = ttnn.linear(hidden_states, self.v_proj_weight, bias=self.v_proj_bias, transpose_b=True)\n",
    "\n",
    "        # Reshape for multi-head attention: split hidden_size into (num_heads * head_dim)\n",
    "        # [sequence_size, batch_size, hidden_size] -> [sequence_size, batch_size * num_heads, head_dim]\n",
    "        q = ttnn.reshape(q, (sequence_size, batch_size * self.num_heads, head_size))\n",
    "        k = ttnn.reshape(k, (sequence_size, batch_size * self.num_heads, head_size))\n",
    "        v = ttnn.reshape(v, (sequence_size, batch_size * self.num_heads, head_size))\n",
    "\n",
    "        # Transpose to bring batch_size * num_heads dimension first for parallel attention computation\n",
    "        # [sequence_size, batch_size * num_heads, head_dim] -> [batch_size * num_heads, sequence_size, head_size]\n",
    "        q = ttnn.transpose(q, 0, 1)\n",
    "        k = ttnn.transpose(k, 0, 1)\n",
    "        v = ttnn.transpose(v, 0, 1)\n",
    "\n",
    "        # Compute attention scores: Q @ K^T\n",
    "        # [batch_size * num_heads, sequence_size, head_size] @ [batch_size * num_heads, head_size, sequence_size] \n",
    "        #   -> [batch_size * num_heads, sequence_size, sequence_size]\n",
    "        scores = ttnn.matmul(q, ttnn.transpose(k, -2, -1))\n",
    "        # Scale scores by 1 / sqrt(head_size) to prevent softmax saturation\n",
    "        scores = scores * self.scale\n",
    "\n",
    "        # Apply attention mask if provided (for causal attention in text encoder)\n",
    "        if self.attention_mask is not None:\n",
    "            # Add mask to scores (mask contains -inf for positions that should be ignored)\n",
    "            # Mask is broadcastable to [batch_size * num_heads, sequence_size, sequence_size]\n",
    "            scores = scores + self.attention_mask\n",
    "\n",
    "        # Apply softmax to get attention weights\n",
    "        # numeric_stable=True uses the numerically stable softmax: softmax(x) = softmax(x - max(x))\n",
    "        # This prevents overflow when computing exp(x) for large values\n",
    "        attn_weights = ttnn.softmax(\n",
    "            scores, dim=-1, numeric_stable=True,\n",
    "        )\n",
    "\n",
    "        # [batch_size * num_heads, sequence_size, sequence_size] @ [batch_size*heads, sequence_size, head_size] \n",
    "        #   -> [batch_size * num_heads, sequence_size, head_size]\n",
    "        attn_output = ttnn.matmul(attn_weights, v)\n",
    "\n",
    "        # Transpose back to sequence-first format\n",
    "        # [batch_size * num_heads, sequence_size, head_size] -> [sequence_size, batch_size * num_heads, head_size]\n",
    "        attn_output = ttnn.transpose(attn_output, 0, 1)\n",
    "\n",
    "        # Merge heads back into hidden dimension\n",
    "        # [sequence_size, batch_size * num_heads, head_size] -> [sequence_size, batch_size, hidden_size]\n",
    "        attn_output = ttnn.reshape(attn_output, (sequence_size, batch_size, hidden_size))\n",
    "\n",
    "        # Apply output projection\n",
    "        dense_out = ttnn.linear(\n",
    "            attn_output,\n",
    "            self.out_proj_weight,\n",
    "            bias=self.out_proj_bias,\n",
    "            transpose_b=True,\n",
    "        )\n",
    "\n",
    "        return dense_out\n",
    "\n",
    "class MultilayerPerceptron:\n",
    "    def __init__(self, state_dict, attention_mask=None, prefix=\"\"):\n",
    "        self.prefix = prefix\n",
    "\n",
    "        self.mlp_c_fc_weight = state_dict[f\"{prefix}.fc1.weight\"]\n",
    "        self.mlp_c_fc_bias = state_dict[f\"{prefix}.fc1.bias\"]\n",
    "        self.mlp_c_proj_weight = state_dict[f\"{prefix}.fc2.weight\"]\n",
    "        self.mlp_c_proj_bias = state_dict[f\"{prefix}.fc2.bias\"]\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = ttnn.linear(x, self.mlp_c_fc_weight, bias=self.mlp_c_fc_bias, transpose_b=True)\n",
    "        x = ttnn.gelu(x)\n",
    "        x = ttnn.linear(x, self.mlp_c_proj_weight, bias=self.mlp_c_proj_bias, transpose_b=True)\n",
    "        return x\n",
    "\n",
    "\n",
    "class ResidualAttentionBlock:\n",
    "    def __init__(self, state_dict, num_heads, attention_mask=None, prefix=\"\"):\n",
    "        self.prefix = prefix\n",
    "        self.num_heads = num_heads\n",
    "\n",
    "        self.attention = MultiHeadAttention(state_dict, num_heads=num_heads, attention_mask=attention_mask, prefix=f\"{prefix}.self_attn\")\n",
    "        self.mlp = MultilayerPerceptron(state_dict, prefix=f\"{prefix}.mlp\")\n",
    "\n",
    "        self.layer_norm_1_weight = state_dict[f\"{prefix}.layer_norm1.weight\"]\n",
    "        self.layer_norm_1_bias = state_dict[f\"{prefix}.layer_norm1.bias\"]\n",
    "        self.layer_norm_2_weight = state_dict[f\"{prefix}.layer_norm2.weight\"]\n",
    "        self.layer_norm_2_bias = state_dict[f\"{prefix}.layer_norm2.bias\"]\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        # LayerNorm\n",
    "        residual = x\n",
    "        x = ttnn.layer_norm(x, weight=self.layer_norm_1_weight, bias=self.layer_norm_1_bias)\n",
    "\n",
    "        # Multihead attention / Self-Attention\n",
    "        # This must be equal to nn.MultiheadAttention(d_model, n_head)(x, x, x, need_weights=False, attn_mask=self.attn_mask)\n",
    "        x = residual + self.attention.forward(x)\n",
    "\n",
    "        # LayerNorm\n",
    "        x_post_layer_norm = ttnn.layer_norm(x, weight=self.layer_norm_2_weight, bias=self.layer_norm_2_bias)\n",
    "\n",
    "        # Multi-Layer Perceptron\n",
    "        x = x + self.mlp.forward(x_post_layer_norm)\n",
    "\n",
    "        return x\n",
    "\n",
    "class Transformer:\n",
    "    def __init__(self, state_dict, num_layers, num_heads, attention_mask=None, prefix=\"\"):\n",
    "        \"\"\"\n",
    "        Initialize a generic Transformer that can be used for both text and vision encoding.\n",
    "        \n",
    "        Args:\n",
    "            state_dict: Model weights dictionary\n",
    "            num_heads: Number of attention heads\n",
    "            attention_mask: Attention mask for causal attention (used for text, None for vision)\n",
    "            prefix: Prefix for layer names in state_dict (e.g., \"text_model.encoder\" or \"vision_model.encoder\")\n",
    "        \"\"\"\n",
    "        self.prefix = prefix\n",
    "        self.layers = []\n",
    "\n",
    "        # Initialize each transformer layer with converted weights\n",
    "        self.layers = [\n",
    "            ResidualAttentionBlock(state_dict, attention_mask=attention_mask, num_heads=num_heads, prefix=f\"{prefix}.layers.{i}\")\n",
    "            for i in range(0, num_layers)\n",
    "        ]\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        for i in range(len(self.layers)):\n",
    "            layer = self.layers[i]\n",
    "            x = layer.forward(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "228fe5e8-5c22-4866-89ce-8651b464df2f",
   "metadata": {},
   "source": [
    "## Vision Transformer Implementation\n",
    "\n",
    "The Vision Transformer class processes images for CLIP. It converts input images into patch embeddings, adds positional encodings, and processes them through transformer layers.\n",
    "\n",
    "### Vision Processing Pipeline\n",
    "\n",
    "1. **Patch Embedding**: Converts 2D image into sequence of patch embeddings using convolution.\n",
    "2. **Class Token**: Prepends a learnable classification token to the sequence.\n",
    "3. **Positional Encoding**: Adds positional information to each patch.\n",
    "4. **Transformer Layers**: Processes the sequence through multiple attention layers.\n",
    "5. **Classification Head**: Extracts features from the class token for final representation.\n",
    "\n",
    "The `forward()` method organizes and executes this entire pipeline, preprocessing image embeddings and calling the generic transformer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe31cab8-b934-4da0-a02c-a2793ba3c989",
   "metadata": {},
   "outputs": [],
   "source": [
    "class VisionTransformer:\n",
    "    def __init__(self, state_dict, num_vision_layers):\n",
    "        self.output_dim = 0\n",
    "\n",
    "        conv2_state_dict_name = \"vision_model.embeddings.patch_embedding.weight\"\n",
    "        self.vision_width = state_dict[conv2_state_dict_name].shape[0]\n",
    "        self.patch_size = state_dict[conv2_state_dict_name].shape[-1]\n",
    "        self.vision_heads = self.vision_width // 64\n",
    "\n",
    "        self.class_embedding = state_dict[\"vision_model.embeddings.class_embedding\"]\n",
    "        self.positional_embedding = state_dict[\"vision_model.embeddings.position_embedding.weight\"]\n",
    "        self.proj = state_dict[\"visual_projection.weight\"]\n",
    "\n",
    "        # Weights preparation for convolution (ttnn.conv2d) must be done on host (CPU)\n",
    "        # To that end, we move convolution weights from device to host and perform its \n",
    "        # layout to Row-Major, which is the preferred layout for TT-NN convolution kernels.\n",
    "        self.conv1_weights = ttnn.from_device(state_dict[conv2_state_dict_name])\n",
    "        self.conv1_weights = ttnn.to_dtype(self.conv1_weights, dtype=ttnn.bfloat16)\n",
    "        self.conv1_weights = ttnn.to_layout(self.conv1_weights, layout=ttnn.ROW_MAJOR_LAYOUT)\n",
    "        \n",
    "        # Layer normalization applied before transformer layers\n",
    "        self.ln_pre_weights = state_dict[\"vision_model.pre_layrnorm.weight\"]\n",
    "        self.ln_pre_bias = state_dict[\"vision_model.pre_layrnorm.bias\"]\n",
    "\n",
    "        # Layer normalization applied after transformer layers (to class token)\n",
    "        self.ln_post_weights = state_dict[\"vision_model.post_layernorm.weight\"]\n",
    "        self.ln_post_bias = state_dict[\"vision_model.post_layernorm.bias\"]\n",
    "\n",
    "        self.transformer = Transformer(\n",
    "            state_dict, num_layers=num_vision_layers, num_heads=self.vision_heads, attention_mask=None, prefix=\"vision_model.encoder\"\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        (batch_size, in_channels, height, width) = x.shape\n",
    "\n",
    "        # === Important: TT-NN conv2d differs from PyTorch in tensor layout ===\n",
    "        # \n",
    "        # PyTorch Conv2d expects: (N, C_in, H, W) - channels-first \"Struct of Arrays\"\n",
    "        # TT-NN conv2d expects:   (N, H, W, C_in) - channels-last \"Array of Structs\"\n",
    "        #\n",
    "        # PyTorch Conv2d output:  (N, C_out, H_out, W_out) - 4D tensor\n",
    "        # TT-NN conv2d output:    (1, 1, N*H_out*W_out, C_out) - flattened 4D tensor\n",
    "        #\n",
    "        # This is why we need to permute and reshape before and after convolution.\n",
    "\n",
    "        # Step 1: Rearrange from channels-first to channels-last layout\n",
    "        # [batch_size, in_channels, height, width] -> [batch_size, height, width, in_channels]\n",
    "        x = ttnn.permute(x, [0, 2, 3, 1])\n",
    "\n",
    "        # Step 2: Convert to row-major layout (required by ttnn.conv2d)\n",
    "        # TT-NN convolution kernels are optimized for row-major data access\n",
    "        x = ttnn.to_layout(x, layout=ttnn.ROW_MAJOR_LAYOUT)\n",
    "\n",
    "        # Output channels for patch embedding (standard ViT uses 768)\n",
    "        out_channels = 768\n",
    "\n",
    "        self.conv1_weights = ttnn.prepare_conv_weights(\n",
    "            weight_tensor=self.conv1_weights,\n",
    "            input_memory_config=x.memory_config(),\n",
    "            input_layout=x.layout,\n",
    "            weights_format=\"OIHW\",\n",
    "            in_channels=in_channels,\n",
    "            out_channels=out_channels,\n",
    "            batch_size=batch_size,\n",
    "            input_height=height,\n",
    "            input_width=width,\n",
    "            kernel_size=[self.patch_size, self.patch_size],\n",
    "            stride=[self.patch_size, self.patch_size],\n",
    "            padding=[0, 0],\n",
    "            dilation=[1, 1],\n",
    "            has_bias=False,\n",
    "            groups=1,\n",
    "            device=get_device(),\n",
    "            input_dtype=x.dtype\n",
    "        )\n",
    "\n",
    "        # Step 3: Apply patch embedding convolution\n",
    "        # This converts the 2D image into a sequence of patch embeddings\n",
    "        # For patch_size=32 and image 224x224: creates (224/32)^2 = 49 patches\n",
    "        x = ttnn.conv2d(\n",
    "            input_tensor=x, \n",
    "            weight_tensor=self.conv1_weights,\n",
    "            in_channels=in_channels,  # Input channels (3 for RGB)\n",
    "            out_channels=out_channels,  # Embedding dimension (768)\n",
    "            batch_size=batch_size,\n",
    "            input_height=height,\n",
    "            input_width=width,\n",
    "            kernel_size=(self.patch_size, self.patch_size),  # Patch size (e.g., 32x32)\n",
    "            stride=(self.patch_size, self.patch_size),  # Non-overlapping patches: stride = kernel_size\n",
    "            padding=(0, 0),  # No padding needed\n",
    "            dilation=(1, 1),  # Standard convolution (no dilation)\n",
    "            groups=0,  # Standard convolution (not grouped/depthwise)\n",
    "            device=get_device(),\n",
    "            return_weights_and_bias=False,  # We already have weights, don't return them\n",
    "            return_output_dim=False,  # We know the output dimensions\n",
    "        )\n",
    "\n",
    "        # Step 4: Reshape convolution output from flattened to sequence format\n",
    "        # Convert to tile layout for subsequent operations (TT-NN's optimized 2D tiled format)\n",
    "        x = ttnn.to_layout(x, layout=ttnn.TILE_LAYOUT)\n",
    "        # Unflatten: [1, 1, N*H_out*W_out, C_out] -> [N, num_patches, embed_dim]\n",
    "        # where num_patches = H_out * W_out = (224/32)^2 = 49\n",
    "        \n",
    "        x = ttnn.reshape(x, (x.shape[0], x.shape[1] * x.shape[2], x.shape[3]))\n",
    "\n",
    "        # Step 5: Prepare the [CLS] token (class embedding)\n",
    "        # ViT prepends a learnable class token to the sequence of patch embeddings\n",
    "        # Reshape class token: [embed_dim] -> [batch_size, 1, embed_dim]\n",
    "        class_embedding = ttnn.reshape(self.class_embedding, (x.shape[0], 1, x.shape[-1]))\n",
    "\n",
    "\n",
    "        # Step 6: Prepare tensors for concatenation\n",
    "        # Move to DRAM memory (slower but more capacity than L1) for concatenation operation\n",
    "        # Note: Concatenation currently requires DRAM memory; future optimizations may use L1 sharded memory\n",
    "        x = ttnn.to_memory_config(x, memory_config=ttnn.DRAM_MEMORY_CONFIG)\n",
    "\n",
    "        # Ensure class_embedding has matching memory configuration\n",
    "        class_embedding = ttnn.to_memory_config(class_embedding, memory_config=x.memory_config())\n",
    "\n",
    "        # Step 7: Prepend class token to patch sequence\n",
    "        # [batch, 1, embed] + [batch, num_patches, embed] -> [batch, num_patches+1, embed]\n",
    "        # This creates: [[CLS], patch_1, patch_2, ..., patch_49]\n",
    "        x = ttnn.concat([class_embedding, x], dim=1, memory_config=None)\n",
    "\n",
    "        # Step 8: Add positional embeddings\n",
    "        # Positional embeddings encode the position of each token in the sequence\n",
    "        # [1, num_patches+1, embed] -> broadcast and add to [batch, num_patches+1, embed]\n",
    "        positional_embedding = ttnn.reshape(self.positional_embedding, (1, x.shape[1], x.shape[2]))\n",
    "        x = x + positional_embedding\n",
    "\n",
    "        # Step 9: Pre-transformer layer normalization\n",
    "        x = ttnn.layer_norm(x, weight=self.ln_pre_weights, bias=self.ln_pre_bias)\n",
    "\n",
    "        # Step 10: Permute to sequence-first format for transformer\n",
    "        # Transformers typically process in sequence-first format\n",
    "        # [batch_size, sequence_size, hidden_size] -> [sequence_size, batch_size, hidden_size]\n",
    "        x = ttnn.permute(x, (1, 0, 2))\n",
    "\n",
    "        # Step 11: Pass through transformer encoder layers\n",
    "        x = self.transformer.forward(x)\n",
    "\n",
    "        # Step 12: Permute back to batch-first format\n",
    "        # [sequence_size, batch_size, hidden_size] -> [batch_size, sequence_size, hidden_size]\n",
    "        x = ttnn.permute(x, (1, 0, 2))\n",
    "\n",
    "        # Step 13: Extract [CLS] token and apply post-layer normalization\n",
    "        # In ViT, the [CLS] token (first token) is used for classification\n",
    "        # x[:, 0, :] extracts the [CLS] token: [batch, seq_len, embed] -> [batch, embed]\n",
    "        x = ttnn.layer_norm(x[:, 0, :], weight=self.ln_post_weights, bias=self.ln_post_bias)\n",
    "\n",
    "        # Step 14: Project to final embedding space (optional projection layer)\n",
    "        # Maps from hidden dimension to the shared vision-language embedding space\n",
    "        if self.proj is not None:\n",
    "            x = ttnn.matmul(x, self.proj, transpose_b=True)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d02ea07a-dc5e-42cd-b214-cdcbfdcacf89",
   "metadata": {},
   "source": [
    "## Complete CLIP Model Implementation\n",
    "\n",
    "We can now define the main CLIP class that combines both text and vision processing capabilities. This class organizes and executes the entire multimodal inference pipeline.\n",
    "\n",
    "### CLIP Architecture Components\n",
    "\n",
    "The CLIP class instantiates and manages the following components:\n",
    "- **Text Transformer**: Processes tokenized text inputs using causal attention masking.\n",
    "- **Vision Transformer**: Processes image inputs through patch-based attention.\n",
    "- **Shared Embedding Space**: Projects both modalities into a common feature space for comparison.\n",
    "\n",
    "### Key Methods\n",
    "- `encode_text()`: Converts text tokens to feature embeddings.\n",
    "- `encode_image()`: Converts images to feature embeddings.  \n",
    "- `forward()`: Performs complete inference and computing similarity scores between images and text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c1ce4f9-c677-4f38-9db3-f066c4da9752",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CLIP:\n",
    "    def __init__(self, state_dict):\n",
    "        self.token_embedding = state_dict[\"text_model.embeddings.token_embedding.weight\"]\n",
    "        self.positional_embedding = state_dict[\"text_model.embeddings.position_embedding.weight\"]\n",
    "\n",
    "        self.text_projection = state_dict[\"text_projection.weight\"]\n",
    "        self.context_length = self.positional_embedding.shape[0]\n",
    "        self.vocab_size = self.token_embedding.shape[0]\n",
    "        self.transformer_width = state_dict[\"text_model.final_layer_norm.weight\"].shape[0]\n",
    "        transformer_heads = self.transformer_width // 64\n",
    "\n",
    "        self.ln_final_weights = state_dict[\"text_model.final_layer_norm.weight\"]\n",
    "        self.ln_final_bias = state_dict[\"text_model.final_layer_norm.bias\"]\n",
    "\n",
    "        self.logit_scale = state_dict[\"logit_scale\"].item()\n",
    "\n",
    "        num_vision_layers = 12 # Hardcoded value for CLIP-ViT-base-patch32\n",
    "        self.visual = VisionTransformer(state_dict, num_vision_layers=num_vision_layers)\n",
    "\n",
    "        num_text_layers = 12 # Hardcoded value for CLIP-ViT-base-patch32\n",
    "        self.transformer = Transformer(\n",
    "            state_dict, num_layers=num_text_layers, num_heads=transformer_heads, attention_mask=self.build_attention_mask(), prefix=\"text_model.encoder\"\n",
    "        )\n",
    "        \n",
    "    def build_attention_mask(self):\n",
    "        \"\"\"\n",
    "        Build causal attention mask for text transformer.\n",
    "        \n",
    "        Causal masking ensures each token can only attend to itself and previous tokens,\n",
    "        preventing the model from \"cheating\" by looking at future tokens. This is essential\n",
    "        for autoregressive language modeling.\n",
    "        \n",
    "        Returns:\n",
    "            Upper triangular mask [context_length, context_length] with -inf above diagonal\n",
    "        \"\"\"\n",
    "        # Create a square mask filled with -inf (tokens cannot attend to masked positions)\n",
    "        # Shape: [context_length, context_length]\n",
    "        mask = ttnn.full(\n",
    "            shape=[self.context_length, self.context_length], \n",
    "            fill_value=float(\"-inf\"), \n",
    "            dtype=ttnn.bfloat16, \n",
    "            device=get_device(), \n",
    "            layout=ttnn.TILE_LAYOUT\n",
    "        )\n",
    "        # Keep only upper triangle (excluding diagonal): prevents attending to future tokens\n",
    "        # diagonal=1 means the diagonal itself is not masked (tokens can attend to themselves)\n",
    "        mask = ttnn.triu(mask, diagonal=1)\n",
    "        return mask\n",
    "\n",
    "\n",
    "    def encode_image(self, image):\n",
    "        return self.visual.forward(image)\n",
    "\n",
    "    def encode_text(self, tokens):\n",
    "        \"\"\"\n",
    "        Encode text tokens into feature embeddings.\n",
    "        \n",
    "        Args:\n",
    "            tokens: Tokenized text input [batch_size, context_length]\n",
    "        \n",
    "        Returns:\n",
    "            Text embeddings in shared vision-language space [batch_size, embed_dim]\n",
    "        \"\"\"\n",
    "        # Convert token IDs to uint32 for embedding lookup\n",
    "        tokens = ttnn.typecast(tokens, dtype=ttnn.uint32)\n",
    "\n",
    "        # Token embedding: [batch, seq_len] -> [batch, seq_len, embed_dim]\n",
    "        x = ttnn.embedding(tokens, weight=self.token_embedding, dtype=ttnn.bfloat16)\n",
    "\n",
    "        # Add learned positional embeddings\n",
    "        # Positional embeddings help the model understand token order\n",
    "        x = x + self.positional_embedding\n",
    "\n",
    "        # Permute to sequence-first format for transformer\n",
    "        # [batch, seq_len, embed] -> [seq_len, batch, embed]\n",
    "        x = ttnn.permute(x, (1, 0, 2))\n",
    "\n",
    "        # Pass through text transformer with causal masking\n",
    "        # Causal masking prevents tokens from attending to future tokens\n",
    "        x = self.transformer.forward(x)\n",
    "\n",
    "        # Permute back to batch-first format\n",
    "        # [seq_len, batch, embed] -> [batch, seq_len, embed]\n",
    "        x = ttnn.permute(x, (1, 0, 2))\n",
    "\n",
    "        # Final layer normalization\n",
    "        x = ttnn.layer_norm(x, weight=self.ln_final_weights, bias=self.ln_final_bias)\n",
    "\n",
    "\n",
    "        # Text Transformer is auto-regressive. This means that the last token has access to all the information in the sequence.\n",
    "        # We can thus extract text features from the end-of-text (EOT) token position\n",
    "        # Note: Using PyTorch for argmax since TT-NN doesn't support advanced indexing yet\n",
    "        torch_tokens = ttnn.to_torch(tokens)\n",
    "        torch_x = ttnn.to_torch(x)\n",
    "\n",
    "        eot_indices = torch_tokens.argmax(dim=-1)  # [batch_size]\n",
    "        torch_selected_features = torch_x[torch.arange(torch_x.shape[0]), eot_indices]  # [batch_size, embed_dim]\n",
    "        \n",
    "        # Move back to TT device and apply text projection\n",
    "        # Projects from transformer hidden size to shared embedding space\n",
    "        x = ttnn.from_torch(torch_selected_features, device=get_device(), layout=ttnn.TILE_LAYOUT)\n",
    "        x = ttnn.matmul(x, self.text_projection, transpose_b=True)\n",
    "        \n",
    "        return x\n",
    "\n",
    "    def forward(self, image, tokens):\n",
    "        \"\"\"\n",
    "        Compute similarity scores between images and text descriptions.\n",
    "        \n",
    "        Args:\n",
    "            image: Preprocessed image tensor [batch_size, channels, height, width]\n",
    "            tokens: Tokenized text tensor [batch_size, context_length]\n",
    "        \n",
    "        Returns:\n",
    "            logits_per_image: Image-to-text similarity scores [batch_size_image, batch_size_text]\n",
    "            logits_per_text: Text-to-image similarity scores [batch_size_text, batch_size_image]\n",
    "        \"\"\"\n",
    "        # Encode both modalities into the shared embedding space\n",
    "        text_features = self.encode_text(tokens)  # [batch_text, embed_dim]\n",
    "        image_features = self.encode_image(image)  # [batch_image, embed_dim]\n",
    "\n",
    "        # Normalize features to unit vectors for cosine similarity\n",
    "        # L2 norm: ||x||_2 = sqrt(sum(x_i^2))\n",
    "        norm_image_features = ttnn.operations.moreh.norm(image_features, p=2.0, dim=1, keepdim=True)\n",
    "        norm_text_features = ttnn.operations.moreh.norm(text_features, p=2.0, dim=1, keepdim=True)\n",
    "\n",
    "        # Normalize: x / ||x|| -> unit vector\n",
    "        image_features = ttnn.divide(image_features, norm_image_features)\n",
    "        text_features = ttnn.divide(text_features, norm_text_features)\n",
    "\n",
    "        # Compute cosine similarity scaled by learned temperature parameter\n",
    "        # logit_scale is learned during training to control the sharpness of the distribution\n",
    "        logit_scale = math.exp(self.logit_scale)\n",
    "\n",
    "        # Compute similarity matrix: scaled dot product of normalized features\n",
    "        # Result: [batch_image, embed] @ [embed, batch_text] = [batch_image, batch_text]\n",
    "        logits_per_image = ttnn.matmul(logit_scale * image_features, text_features, transpose_b=True)\n",
    "        # Transpose for text-to-image direction\n",
    "        logits_per_text = ttnn.transpose(logits_per_image, 0, 1)\n",
    "\n",
    "        return logits_per_image, logits_per_text\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "940036cb-a7e8-43ce-9716-36a84d362d2e",
   "metadata": {},
   "source": [
    "## Image Preprocessing\n",
    "\n",
    "While input images can have any dimensions and color spaces, CLIP expects standardized 224×224 RGB images. Images are pre-processed to match the model's expected input format.\n",
    "\n",
    "### Preprocessing Pipeline\n",
    "\n",
    "Preprocessing applies the following transformations in sequence:\n",
    "1. **Resize**: Scale image to 224×224 pixels using bicubic interpolation.\n",
    "2. **Center Crop**: Crop the center region to ensure exact dimensions.\n",
    "3. **RGB Conversion**: Convert to RGB color space if needed.\n",
    "4. **Normalization**: Apply ImageNet normalization statistics used during CLIP training.\n",
    "\n",
    "Preprocessing ensures consistent input format regardless of the original image properties. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e08e26c-5fe8-4162-9cf7-37e2c5ce24f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_image(image, model_resolution):\n",
    "    def _convert_image_to_rgb(image):\n",
    "        return image.convert(\"RGB\")\n",
    "\n",
    "    # Pre-process image on host with torch\n",
    "    transform_fn = Compose(\n",
    "        [\n",
    "            Resize(model_resolution, interpolation=InterpolationMode.BICUBIC),\n",
    "            CenterCrop(model_resolution),\n",
    "            _convert_image_to_rgb,\n",
    "            ToTensor(),\n",
    "            Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),\n",
    "        ]\n",
    "    )\n",
    "    return transform_fn(image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be67cd47-2f80-490b-98a1-aa6a7ba84832",
   "metadata": {},
   "source": [
    "## Image Download Utility\n",
    "\n",
    "We use a utility function to download images from URLs for demonstrations. This function handles HTTP requests and converts the response into a PIL Image object that can be processed by the preprocessing pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8108429-2409-4a3c-9a77-c0adbb4537ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "def download_image(url):\n",
    "    \"\"\"\n",
    "    Download an image from a URL and return it as a PIL Image object.\n",
    "    \n",
    "    Args:\n",
    "        url (str): The URL of the image to download\n",
    "        \n",
    "    Returns:\n",
    "        PIL.Image: The downloaded image\n",
    "    \"\"\"\n",
    "    try:\n",
    "        response = requests.get(url, timeout=30)\n",
    "        response.raise_for_status()  # Raise an exception for bad status codes\n",
    "        \n",
    "        # Convert the response content to a PIL Image\n",
    "        image = Image.open(BytesIO(response.content))\n",
    "        return image\n",
    "    except requests.RequestException as e:\n",
    "        raise Exception(f\"Failed to download image from {url}: {e}\")\n",
    "    except Exception as e:\n",
    "        raise Exception(f\"Failed to process downloaded image: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61b62fd9",
   "metadata": {},
   "source": [
    "## Model and Tokenizer downloading\n",
    "\n",
    "We define two functions to download the weights and tokenizer from Hugging Face.\n",
    "\n",
    "For practical purposes, we can also specify a `TTNN_TUTORIALS_MODELS_CLIP_PATH` environment variable to avoid downloading the model.\n",
    "If it is defined, then model and weights will be loaded from the location indicated by `TTNN_TUTORIALS_MODELS_CLIP_PATH`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3966d496",
   "metadata": {},
   "outputs": [],
   "source": [
    "def download_model(model_name):\n",
    "    clip_model_location = model_name # By default, download from Hugging Face\n",
    "    \n",
    "    # If TTNN_TUTORIALS_MODELS_CLIP_PATH is set, use it as the cache directory to avoid requests to Hugging Face\n",
    "    cache_dir = os.getenv(\"TTNN_TUTORIALS_MODELS_CLIP_PATH\")\n",
    "    if cache_dir is not None:\n",
    "        clip_model_location = cache_dir\n",
    "\n",
    "    # Load model weights (download if cache_dir was not set)\n",
    "    model = CLIPModel.from_pretrained(clip_model_location)\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "def download_tokenizer(tokenizer_name):\n",
    "    clip_tokenizer_location = tokenizer_name # By default, download from Hugging Face\n",
    "\n",
    "    # If TTNN_TUTORIALS_MODELS_CLIP_PATH is set, use it as the cache directory to avoid requests to Hugging Face\n",
    "    cache_dir = os.getenv(\"TTNN_TUTORIALS_MODELS_CLIP_PATH\")\n",
    "    if cache_dir is not None:\n",
    "        clip_tokenizer_location = cache_dir\n",
    "\n",
    "    tokenizer = CLIPTokenizer.from_pretrained(clip_tokenizer_location)\n",
    "\n",
    "    return tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfa56ff5-8698-4705-8aee-32f81f61156b",
   "metadata": {},
   "source": [
    "## Running CLIP Inference\n",
    "\n",
    "Having defined each component in the CLIP model, we can perform inference on an input image and text prompts. This section demonstrates the complete inference pipeline from loading pre-trained weights to computing similarity scores.\n",
    "\n",
    "### Inference Pipeline:\n",
    "\n",
    "1. **Device Initialization**: Open TT-NN device with configured L1 cache size.\n",
    "2. **Model Loading**: Download pre-trained CLIP weights from Hugging Face using `CLIPModel.from_pretrained()`.\n",
    "3. **Weight Conversion**: Convert PyTorch weights to TT-NN format for hardware acceleration.\n",
    "4. **Image Processing**: \n",
    "   - Download image from URL.\n",
    "   - Preprocess (resize to 224×224, normalize with ImageNet statistics)\n",
    "   - Convert to TT-NN tensor with bfloat16 precision.\n",
    "5. **Text Processing**: \n",
    "   - Tokenize text prompts using `CLIPTokenizer`.\n",
    "   - Add padding to match context length.\n",
    "   - Convert to TT-NN tensors.\n",
    "6. **Forward Pass**: \n",
    "   - Encode image through Vision Transformer.\n",
    "   - Encode text through Text Transformer.\n",
    "   - Normalize both embeddings to unit vectors.\n",
    "   - Compute cosine similarity (dot product of normalized vectors).\n",
    "7. **Results**: Apply softmax to convert similarity scores to probabilities.\n",
    "\n",
    "### Text Tokenization\n",
    "\n",
    "Since TT-NN does not handle tokenization natively, we use the `CLIPTokenizer` from the `transformers` library. The tokenizer performs the following tasks:\n",
    "- Converts text strings into token IDs matching CLIP's vocabulary (49,408 tokens).\n",
    "- Adds special tokens: [SOS] (start of sequence) and [EOS] (end of sequence).\n",
    "- Pads sequences to the model's context length (77 tokens for CLIP).\n",
    "- The [EOS] token position is used to extract the final text representation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64691ea0-cbf8-452c-b9d3-4b57157a480a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize TT-NN device for hardware acceleration\n",
    "open_ttnn()\n",
    "\n",
    "\n",
    "# Load pre-trained CLIP model and convert weights to TT-NN format\n",
    "logger.info(\"Loading pre-trained CLIP model...\")\n",
    "\n",
    "model = download_model(\"openai/clip-vit-base-patch32\")\n",
    "state_dict = convert_model_to_ttnn(model.state_dict())\n",
    "\n",
    "tokenizer = download_tokenizer(\"openai/clip-vit-base-patch32\")\n",
    "\n",
    "\n",
    "# Initialize our TT-NN CLIP implementation\n",
    "clip = CLIP(state_dict)\n",
    "\n",
    "# Download and preprocess test image\n",
    "logger.info(\"Downloading and preprocessing image...\")\n",
    "image_url = \"https://media.githubusercontent.com/media/tenstorrent/tutorial-assets/refs/heads/main/media/clip_tutorial/CLIP.png\"\n",
    "image = download_image(image_url)\n",
    "\n",
    "# Preprocess image to model requirements (224x224, normalized with ImageNet statistics)\n",
    "# unsqueeze(0) adds batch dimension: [C, H, W] -> [1, C, H, W]\n",
    "image = preprocess_image(image, 224).unsqueeze(0).to(\"cpu\")\n",
    "\n",
    "# Convert PyTorch image tensor to TT-NN tensor with bfloat16 precision\n",
    "# bfloat16 provides good balance between precision and memory/compute efficiency\n",
    "preferred_dtype = ttnn.bfloat16\n",
    "tt_image = ttnn.from_torch(image, device=get_device(), layout=ttnn.TILE_LAYOUT, dtype=preferred_dtype)\n",
    "\n",
    "# Define text prompts for zero-shot classification\n",
    "# The model will compute similarity between the image and each text description\n",
    "prompts = [\"a diagram\", \"a dog\", \"a cat\"]\n",
    "\n",
    "# Tokenize text prompts using CLIP's tokenizer\n",
    "logger.info(\"Tokenizing text prompts...\")\n",
    "# padding=\"max_length\" ensures all sequences are padded to context_length (77 tokens)\n",
    "# return_tensors=\"pt\" returns PyTorch tensors\n",
    "tokenized_inputs = tokenizer(prompts, padding=\"max_length\", max_length=clip.context_length, return_tensors=\"pt\")\n",
    "tokens_pretrained_host = tokenized_inputs[\"input_ids\"]  # Shape: [num_prompts, context_length]\n",
    "# Convert tokenized text to TT-NN tensors for device execution\n",
    "tokens_pretrained = ttnn.from_torch(tokens_pretrained_host, device=get_device(), layout=ttnn.TILE_LAYOUT)\n",
    "\n",
    "\n",
    "# Perform CLIP inference: compute similarity between image and text\n",
    "logger.info(\"Running CLIP inference...\")\n",
    "time_start = time.time()\n",
    "logits_per_image, logits_per_text = clip.forward(tt_image, tokens_pretrained)\n",
    "time_end = time.time()\n",
    "logger.info(f\"Time taken: {time_end - time_start:.3f} seconds\")\n",
    "\n",
    "# Convert logits (similarity scores) to probabilities using softmax\n",
    "# Softmax normalizes scores so they sum to 1.0, representing a probability distribution\n",
    "probs = ttnn.softmax(logits_per_image, dim=-1)\n",
    "logger.info(f\"==== Zero-shot Classification Results ====\")\n",
    "logger.info(f\"Image: {image_url.split('/')[-1]}\")\n",
    "logger.info(f\"Classification probabilities:\")\n",
    "\n",
    "# Display results sorted by probability (highest first)\n",
    "probs_torch = ttnn.to_torch(probs)\n",
    "results = [(prompt, probs_torch[0][i].item()) for i, prompt in enumerate(prompts)]\n",
    "results.sort(key=lambda x: x[1], reverse=True)\n",
    "\n",
    "for prompt, prob in results:\n",
    "    logger.info(f\"  '{prompt}': {prob:.4f} ({prob*100:.2f}%)\")\n",
    "\n",
    "# Clean up resources\n",
    "close_ttnn()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
